import matplotlib.pyplot as plt
import numpy as np

# 读取数据文件
def parse_loss_data(file_path):
    epochs = []
    losses = []
    with open(file_path, 'r') as f:
        for line in f:
            print(line)
            if 'Epoch' in line and 'Loss' in line:
                parts = line.split(',')
                print(parts)
                epoch_part = parts[0].strip().split('[')[1].split('/')[0]
                loss_part = parts[1].strip().split(':')[1]
                epochs.append(int(epoch_part))
                losses.append(float(loss_part))
    return np.array(epochs), np.array(losses)

# 绘制损失曲线
def plot_loss(epochs, losses):
    plt.figure(figsize=(10, 6))
    
    # 使用对数y轴使下降趋势更明显
    plt.yscale('log')
    
    plt.plot(epochs, losses, 'b-', linewidth=1)
    plt.title('Training Loss Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Loss (log scale)')
    plt.grid(True, which="both", ls="--")
    
    # 突出显示用户关注的区域
    highlight_start = 6570
    highlight_end = 6690
    mask = (epochs >= highlight_start) & (epochs <= highlight_end)
    plt.plot(epochs[mask], losses[mask], 'r-', linewidth=2)
    
    plt.show()

if __name__ == '__main__':
    file_path = 'data.txt'
    epochs, losses = parse_loss_data(file_path)
    plot_loss(epochs, losses)
    # Test Loss: 0.0719
    # Epoch [10000/10000], Loss: 0.0133
    # Test Loss: 0.0722
    # Epoch [10000/10000], Loss: 0.0152

